# -*- coding: utf-8 -*-
import torch


class TrainConfig:
    def __init__(self):
        # 词典路径：在vocab文件夹里面
        self.tokenizer_path = 'vocab/vocab.txt'

        # 训练文件路径
        self.train_path = 'data/medical_train.pkl'

        # 验证数据文件路径
        self.valid_path = 'data/medical_valid.pkl'

        # 模型配置文件
        self.model_config = 'model_config/gpt2_config.json'

        # 模型保存路径
        self.save_model_path = 'save_model'

        # 预训练模型路径
        self.pretrained_model = ''

        # 训练轮次
        self.epochs = 400

        # 训练批次
        self.batch_size = 8

        # 优化器 学习率 梯度下降
        self.lr = 2.6e-5
        # 优化器 为了增加数值计算的稳定性而加到分母里的项，其为了防止在实现中除以零
        self.eps = 1.0e-09

        # 调度器 梯度累积步长 防止批次较小引起的不稳定
        self.gradient_accumulation_steps = 4
        # 调度器 默认 4000
        self.warmup_steps = 100
        # 使用Warmup预热学习率的方式,即先用最初的小学习率训练，然后每个step增大一点点，
        # 直到达到最初设置的比较大的学习率时（注：此时预热学习率完成），
        # 采用最初设置的学习率进行训练（注：预热学习率完成后的训练过程，学习率是衰减的），有助于使模型收敛速度变快，效果更佳。

        # 允许的最大梯度范数（阈值）
        self.max_grad_norm = 2.0

        # 多少步汇报一次loss
        self.loss_step = 1

        # 训练几轮保存一次模型
        self.save_model_epochs = 10

        # 忽略一些字符：句子需要长度补齐，针对补的部分，没有意义，所以一般不进行梯度更新
        self.ignore_index = -100


class InferenceConfig:
    def __init__(self):
        # 词典路径：在vocab文件夹里面
        self.tokenizer_path = 'vocab/vocab.txt'

        # 使用模型路径
        self.used_model_path = 'save_model/bj_epoch140'

        # 历史对话句子的长度
        self.max_history_len = 2

        # 每一个完整对话的句子最大长度
        # '每个utterance的最大长度,超过指定长度则进行截断,默认25'
        self.max_token = 300

        # "重复惩罚参数，若生成的对话重复性较高，可适当提高该参数"
        self.repetition_penalty = 10

        # 输出的最高单词数量
        self.top_k = 4  # '最高k选1。默认8'

        # 输出的最高几个单词概率和
        self.top_p = 0.9


if __name__ == '__main__':
    pc = TrainConfig()
    print(pc.train_path)
    print(torch.cuda.device_count())
